import os

import numpy as np
import pytest
import torch
from numpy.testing import assert_almost_equal, assert_array_almost_equal
from PIL import Image
from ultralytics.engine.results import Results

from datachain.model.ultralytics import (
    YoloBBox,
    YoloBBoxes,
    YoloOBBox,
    YoloOBBoxes,
    YoloPose,
    YoloPoses,
    YoloSegment,
    YoloSegments,
)


@pytest.fixture
def running_img() -> np.ndarray:
    img_file = os.path.join(os.path.dirname(__file__), "data", "running.jpg")
    with Image.open(img_file) as img:
        return np.array(img)


@pytest.fixture
def ships_img() -> np.ndarray:
    img_file = os.path.join(os.path.dirname(__file__), "data", "ships.jpg")
    with Image.open(img_file) as img:
        return np.array(img)


@pytest.fixture
def running_img_masks() -> torch.Tensor:
    mask0_file = os.path.join(os.path.dirname(__file__), "data", "running-mask0.png")
    with Image.open(mask0_file) as mask0_img:
        mask0_np = np.array(mask0_img)

    mask1_file = os.path.join(os.path.dirname(__file__), "data", "running-mask1.png")
    with Image.open(mask1_file) as mask1_img:
        mask1_np = np.array(mask1_img)

    stacked = np.stack(
        [mask0_np.astype(np.float32), mask1_np.astype(np.float32)],
        axis=0,
    )
    return torch.from_numpy(stacked)


def test_yolo_bbox_from_results_empty(running_img):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={},
        boxes=torch.empty((0, 6)),
    )

    model = YoloBBox.from_result(result)
    assert model.model_dump() == {
        "cls": -1,
        "name": "",
        "confidence": 0,
        "box": {
            "coords": [],
            "title": "",
        },
    }

    model = YoloBBoxes.from_results([result])
    assert model.model_dump() == {
        "cls": [],
        "name": [],
        "confidence": [],
        "box": [],
    }


def test_yolo_bbox_from_results(running_img):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={0: "person", 16: "dog"},
        boxes=torch.tensor(
            [
                [100.2483, 83.7399, 183.179, 238.1918, 0.9057, 0.0],
                [10.8968, 177.1552, 71.1275, 239.0617, 0.8919, 16.0],
            ],
        ),
    )

    model = YoloBBox.from_result(result)
    model_json = model.model_dump()
    assert set(model_json.keys()) == {"cls", "name", "confidence", "box"}
    assert model_json["cls"] == 0
    assert model_json["name"] == "person"
    assert_almost_equal(model_json["confidence"], 0.9057, decimal=3)
    assert model_json["box"] == {"coords": [100, 84, 183, 238], "title": "person"}

    model = YoloBBoxes.from_results([result])
    model_json = model.model_dump()
    assert set(model_json.keys()) == {"cls", "name", "confidence", "box"}
    assert model_json["cls"] == [0, 16]
    assert model_json["name"] == ["person", "dog"]
    assert_array_almost_equal(model_json["confidence"], [0.9057, 0.8919], decimal=3)
    assert model_json["box"] == [
        {"coords": [100, 84, 183, 238], "title": "person"},
        {"coords": [11, 177, 71, 239], "title": "dog"},
    ]


def test_yolo_obbox_from_results_empty(ships_img):
    result = Results(
        orig_img=ships_img,
        path="ships.jpeg",
        names={},
        obb=torch.empty((0, 7)),
    )

    model = YoloOBBox.from_result(result)
    assert model.model_dump() == {
        "cls": -1,
        "name": "",
        "confidence": 0,
        "box": {
            "coords": [],
            "title": "",
        },
    }

    model = YoloOBBoxes.from_results([result])
    assert model.model_dump() == {
        "cls": [],
        "name": [],
        "confidence": [],
        "box": [],
    }


def test_yolo_obbox_from_results(ships_img):
    result = Results(
        orig_img=ships_img,
        path="ships.jpeg",
        names={1: "ship"},
        obb=torch.tensor(
            [
                [272.7724, 83.4637, 80.2673, 41.5667, 0.5629, 0.7289, 1.0],
                [70.2, 95.1912, 143.2061, 48.404, 0.494, 0.5785, 1.0],
                [120.6607, 76.8303, 115.6623, 42.3174, 0.7043, 0.5677, 1.0],
            ]
        ),
    )

    model = YoloOBBox.from_result(result)
    model_dict = model.model_dump()
    assert set(model_dict.keys()) == {"cls", "name", "confidence", "box"}
    assert model_dict["cls"] == 1
    assert model_dict["name"] == "ship"
    assert_almost_equal(model_dict["confidence"], 0.7289, decimal=3)
    assert model_dict["box"] == {
        "coords": [296, 122, 318, 87, 250, 44, 228, 80],
        "title": "ship",
    }

    model = YoloOBBoxes.from_results([result])
    model_dict = model.model_dump()
    assert set(model_dict.keys()) == {"cls", "name", "confidence", "box"}
    assert model_dict["cls"] == [1, 1, 1]
    assert model_dict["name"] == ["ship", "ship", "ship"]
    assert_array_almost_equal(
        model_dict["confidence"], [0.7289, 0.5785, 0.5677], decimal=3
    )
    assert model_dict["box"] == [
        {"coords": [296, 122, 318, 87, 250, 44, 228, 80], "title": "ship"},
        {"coords": [122, 150, 145, 108, 19, 40, -4, 83], "title": "ship"},
        {"coords": [151, 130, 178, 98, 90, 23, 63, 56], "title": "ship"},
    ]


def test_yolo_segment_from_results_empty(running_img):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={},
        obb=torch.empty((0, 7)),
    )

    model = YoloSegment.from_result(result)
    assert model.model_dump() == {
        "cls": -1,
        "name": "",
        "confidence": 0,
        "box": {
            "coords": [],
            "title": "",
        },
        "segment": {
            "title": "",
            "x": [],
            "y": [],
        },
    }

    model = YoloSegments.from_results([result])
    assert model.model_dump() == {
        "cls": [],
        "name": [],
        "confidence": [],
        "box": [],
        "segment": [],
    }


def test_yolo_segment_from_results_empty_segments(running_img):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={0: "person"},
        boxes=torch.tensor([[102.0, 84.0, 183.0, 238.0, 0.9078, 0.0]]),
    )

    model = YoloSegment.from_result(result)
    assert model.model_dump() == {
        "cls": 0,
        "name": "person",
        "confidence": 0.9078,
        "box": {
            "coords": [102, 84, 183, 238],
            "title": "person",
        },
        "segment": {
            "title": "",
            "x": [],
            "y": [],
        },
    }

    model = YoloSegments.from_results([result])
    assert model.model_dump() == {
        "cls": [0],
        "name": ["person"],
        "confidence": [0.9078],
        "box": [
            {
                "coords": [102, 84, 183, 238],
                "title": "person",
            }
        ],
        "segment": [],
    }


def test_yolo_seg_from_results(running_img, running_img_masks):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={0: "person", 16: "dog"},
        boxes=torch.tensor(
            [
                [11.6684, 178.2727, 71.4605, 238.8026, 0.9157, 16.0],
                [100.8046, 84.3006, 183.2971, 238.1916, 0.8578, 0.0],
            ],
        ),
        masks=running_img_masks,
    )

    model = YoloSegment.from_result(result)
    model_dict = model.model_dump()
    assert set(model_dict.keys()) == {
        "cls",
        "name",
        "confidence",
        "box",
        "segment",
    }
    assert model_dict["cls"] == 16
    assert model_dict["name"] == "dog"
    assert_almost_equal(model_dict["confidence"], 0.9157, decimal=3)
    assert model_dict["box"] == {
        "coords": [12, 178, 71, 239],
        "title": "dog",
    }
    assert model_dict["segment"] == {
        "title": "dog",
        "x": [
            16,
            16,
            16,
            16,
            15,
            15,
            15,
            15,
            14,
            14,
            14,
            14,
            13,
            13,
            12,
            11,
            11,
            13,
            14,
            14,
            16,
            16,
            17,
            17,
            20,
            20,
            21,
            21,
            21,
            21,
            22,
            22,
            22,
            22,
            23,
            23,
            22,
            22,
            22,
            22,
            21,
            21,
            21,
            21,
            20,
            20,
            20,
            20,
            19,
            19,
            19,
            19,
            18,
            18,
            17,
            17,
            16,
            16,
            16,
            17,
            18,
            23,
            24,
            25,
            26,
            26,
            26,
            26,
            27,
            27,
            29,
            29,
            30,
            31,
            31,
            32,
            32,
            36,
            36,
            38,
            38,
            39,
            39,
            40,
            40,
            40,
            40,
            41,
            41,
            41,
            41,
            43,
            43,
            45,
            45,
            46,
            47,
            47,
            50,
            50,
            51,
            52,
            52,
            53,
            54,
            56,
            56,
            59,
            59,
            60,
            60,
            61,
            61,
            61,
            61,
            62,
            62,
            62,
            62,
            62,
            62,
            63,
            63,
            65,
            66,
            66,
            67,
            67,
            68,
            68,
            68,
            68,
            69,
            69,
            68,
            68,
            69,
            69,
            69,
            69,
            70,
            70,
            70,
            70,
            70,
            70,
            71,
            72,
            72,
            71,
            70,
            70,
            70,
            70,
            70,
            70,
            69,
            69,
            69,
            69,
            68,
            68,
            67,
            67,
            67,
            67,
            66,
            66,
            65,
            65,
            65,
            65,
            64,
            64,
            63,
            63,
            62,
            62,
            62,
            61,
            61,
            60,
            59,
            59,
            58,
            58,
            57,
            57,
            56,
            56,
            55,
            55,
            55,
            54,
            52,
            52,
            51,
            51,
            49,
            49,
            48,
            44,
            44,
            42,
            41,
            39,
            38,
            37,
            36,
            36,
            35,
            35,
            34,
            34,
            33,
            33,
            32,
            32,
            31,
            31,
            30,
            30,
            25,
            25,
            24,
            23,
            21,
            21,
            20,
            19,
            19,
            18,
            18,
            18,
        ],
        "y": [
            179,
            180,
            181,
            181,
            182,
            182,
            183,
            185,
            185,
            186,
            186,
            187,
            187,
            187,
            188,
            188,
            200,
            200,
            201,
            201,
            202,
            202,
            203,
            204,
            207,
            208,
            208,
            209,
            210,
            211,
            211,
            212,
            212,
            213,
            214,
            219,
            220,
            222,
            222,
            223,
            224,
            225,
            226,
            227,
            227,
            228,
            228,
            229,
            229,
            230,
            230,
            231,
            231,
            232,
            233,
            234,
            235,
            235,
            236,
            236,
            236,
            236,
            235,
            235,
            236,
            236,
            237,
            237,
            237,
            238,
            238,
            237,
            237,
            237,
            237,
            235,
            235,
            232,
            231,
            229,
            229,
            228,
            227,
            226,
            225,
            225,
            225,
            224,
            224,
            223,
            223,
            221,
            221,
            219,
            219,
            218,
            218,
            218,
            218,
            218,
            218,
            219,
            219,
            219,
            219,
            222,
            223,
            225,
            225,
            227,
            227,
            228,
            228,
            229,
            230,
            230,
            235,
            235,
            236,
            236,
            237,
            237,
            237,
            237,
            237,
            237,
            236,
            236,
            236,
            235,
            235,
            233,
            232,
            229,
            228,
            226,
            225,
            221,
            221,
            219,
            219,
            218,
            217,
            216,
            216,
            215,
            215,
            215,
            202,
            202,
            202,
            201,
            201,
            200,
            199,
            198,
            198,
            197,
            197,
            196,
            196,
            195,
            195,
            194,
            194,
            193,
            192,
            192,
            191,
            190,
            190,
            189,
            189,
            188,
            187,
            187,
            186,
            186,
            186,
            186,
            185,
            185,
            184,
            184,
            184,
            184,
            183,
            183,
            183,
            183,
            182,
            182,
            183,
            183,
            185,
            186,
            187,
            187,
            188,
            188,
            189,
            189,
            188,
            188,
            188,
            188,
            187,
            187,
            187,
            187,
            187,
            187,
            186,
            186,
            186,
            186,
            185,
            185,
            184,
            183,
            183,
            179,
            179,
            180,
            181,
            181,
            183,
            183,
            183,
            183,
            183,
            183,
            182,
            179,
        ],
    }

    model = YoloSegments.from_results([result])
    model_dict = model.model_dump()
    assert set(model_dict.keys()) == {
        "cls",
        "name",
        "confidence",
        "box",
        "segment",
    }
    assert model_dict["cls"] == [16, 0]
    assert model_dict["name"] == ["dog", "person"]
    assert_array_almost_equal(model_dict["confidence"], [0.9157, 0.8578], decimal=3)
    assert model_dict["box"] == [
        {"coords": [12, 178, 71, 239], "title": "dog"},
        {"coords": [101, 84, 183, 238], "title": "person"},
    ]
    assert model_dict["segment"] == [
        {
            "title": "dog",
            "x": [
                16,
                16,
                16,
                16,
                15,
                15,
                15,
                15,
                14,
                14,
                14,
                14,
                13,
                13,
                12,
                11,
                11,
                13,
                14,
                14,
                16,
                16,
                17,
                17,
                20,
                20,
                21,
                21,
                21,
                21,
                22,
                22,
                22,
                22,
                23,
                23,
                22,
                22,
                22,
                22,
                21,
                21,
                21,
                21,
                20,
                20,
                20,
                20,
                19,
                19,
                19,
                19,
                18,
                18,
                17,
                17,
                16,
                16,
                16,
                17,
                18,
                23,
                24,
                25,
                26,
                26,
                26,
                26,
                27,
                27,
                29,
                29,
                30,
                31,
                31,
                32,
                32,
                36,
                36,
                38,
                38,
                39,
                39,
                40,
                40,
                40,
                40,
                41,
                41,
                41,
                41,
                43,
                43,
                45,
                45,
                46,
                47,
                47,
                50,
                50,
                51,
                52,
                52,
                53,
                54,
                56,
                56,
                59,
                59,
                60,
                60,
                61,
                61,
                61,
                61,
                62,
                62,
                62,
                62,
                62,
                62,
                63,
                63,
                65,
                66,
                66,
                67,
                67,
                68,
                68,
                68,
                68,
                69,
                69,
                68,
                68,
                69,
                69,
                69,
                69,
                70,
                70,
                70,
                70,
                70,
                70,
                71,
                72,
                72,
                71,
                70,
                70,
                70,
                70,
                70,
                70,
                69,
                69,
                69,
                69,
                68,
                68,
                67,
                67,
                67,
                67,
                66,
                66,
                65,
                65,
                65,
                65,
                64,
                64,
                63,
                63,
                62,
                62,
                62,
                61,
                61,
                60,
                59,
                59,
                58,
                58,
                57,
                57,
                56,
                56,
                55,
                55,
                55,
                54,
                52,
                52,
                51,
                51,
                49,
                49,
                48,
                44,
                44,
                42,
                41,
                39,
                38,
                37,
                36,
                36,
                35,
                35,
                34,
                34,
                33,
                33,
                32,
                32,
                31,
                31,
                30,
                30,
                25,
                25,
                24,
                23,
                21,
                21,
                20,
                19,
                19,
                18,
                18,
                18,
            ],
            "y": [
                179,
                180,
                181,
                181,
                182,
                182,
                183,
                185,
                185,
                186,
                186,
                187,
                187,
                187,
                188,
                188,
                200,
                200,
                201,
                201,
                202,
                202,
                203,
                204,
                207,
                208,
                208,
                209,
                210,
                211,
                211,
                212,
                212,
                213,
                214,
                219,
                220,
                222,
                222,
                223,
                224,
                225,
                226,
                227,
                227,
                228,
                228,
                229,
                229,
                230,
                230,
                231,
                231,
                232,
                233,
                234,
                235,
                235,
                236,
                236,
                236,
                236,
                235,
                235,
                236,
                236,
                237,
                237,
                237,
                238,
                238,
                237,
                237,
                237,
                237,
                235,
                235,
                232,
                231,
                229,
                229,
                228,
                227,
                226,
                225,
                225,
                225,
                224,
                224,
                223,
                223,
                221,
                221,
                219,
                219,
                218,
                218,
                218,
                218,
                218,
                218,
                219,
                219,
                219,
                219,
                222,
                223,
                225,
                225,
                227,
                227,
                228,
                228,
                229,
                230,
                230,
                235,
                235,
                236,
                236,
                237,
                237,
                237,
                237,
                237,
                237,
                236,
                236,
                236,
                235,
                235,
                233,
                232,
                229,
                228,
                226,
                225,
                221,
                221,
                219,
                219,
                218,
                217,
                216,
                216,
                215,
                215,
                215,
                202,
                202,
                202,
                201,
                201,
                200,
                199,
                198,
                198,
                197,
                197,
                196,
                196,
                195,
                195,
                194,
                194,
                193,
                192,
                192,
                191,
                190,
                190,
                189,
                189,
                188,
                187,
                187,
                186,
                186,
                186,
                186,
                185,
                185,
                184,
                184,
                184,
                184,
                183,
                183,
                183,
                183,
                182,
                182,
                183,
                183,
                185,
                186,
                187,
                187,
                188,
                188,
                189,
                189,
                188,
                188,
                188,
                188,
                187,
                187,
                187,
                187,
                187,
                187,
                186,
                186,
                186,
                186,
                185,
                185,
                184,
                183,
                183,
                179,
                179,
                180,
                181,
                181,
                183,
                183,
                183,
                183,
                183,
                183,
                182,
                179,
            ],
        },
        {
            "title": "person",
            "x": [
                136,
                136,
                132,
                132,
                133,
                133,
                134,
                134,
                134,
                134,
                135,
                135,
                135,
                135,
                136,
                136,
                135,
                135,
                132,
                132,
                130,
                130,
                129,
                129,
                128,
                128,
                128,
                128,
                127,
                127,
                127,
                127,
                126,
                126,
                126,
                126,
                125,
                125,
                125,
                125,
                124,
                124,
                124,
                124,
                123,
                123,
                123,
                123,
                122,
                122,
                122,
                122,
                122,
                122,
                121,
                121,
                121,
                121,
                120,
                120,
                119,
                119,
                118,
                118,
                117,
                117,
                117,
                117,
                116,
                116,
                115,
                115,
                115,
                115,
                115,
                115,
                114,
                114,
                114,
                114,
                113,
                113,
                111,
                110,
                110,
                109,
                108,
                108,
                107,
                107,
                107,
                106,
                106,
                105,
                105,
                104,
                103,
                101,
                101,
                103,
                104,
                104,
                105,
                106,
                107,
                107,
                108,
                109,
                109,
                111,
                112,
                112,
                113,
                113,
                114,
                115,
                115,
                115,
                116,
                116,
                117,
                119,
                120,
                120,
                120,
                120,
                121,
                121,
                122,
                122,
                123,
                124,
                125,
                125,
                125,
                126,
                126,
                128,
                128,
                128,
                129,
                130,
                130,
                130,
                130,
                130,
                130,
                129,
                129,
                129,
                129,
                129,
                129,
                130,
                130,
                130,
                130,
                133,
                133,
                133,
                133,
                134,
                134,
                134,
                134,
                135,
                135,
                135,
                135,
                136,
                136,
                136,
                136,
                137,
                137,
                137,
                137,
                138,
                138,
                138,
                138,
                138,
                138,
                139,
                139,
                139,
                139,
                140,
                140,
                140,
                140,
                140,
                140,
                139,
                139,
                139,
                139,
                138,
                138,
                138,
                138,
                138,
                138,
                137,
                137,
                137,
                137,
                136,
                136,
                135,
                135,
                134,
                133,
                132,
                129,
                129,
                128,
                127,
                127,
                126,
                126,
                127,
                127,
                128,
                128,
                129,
                130,
                131,
                132,
                132,
                133,
                134,
                138,
                138,
                144,
                144,
                145,
                145,
                145,
                145,
                145,
                145,
                144,
                144,
                143,
                143,
                144,
                144,
                144,
                144,
                145,
                145,
                145,
                145,
                145,
                145,
                146,
                146,
                146,
                146,
                147,
                147,
                147,
                147,
                148,
                148,
                148,
                148,
                149,
                149,
                149,
                149,
                150,
                150,
                150,
                150,
                151,
                151,
                151,
                151,
                152,
                152,
                153,
                159,
                159,
                160,
                161,
                161,
                162,
                162,
                161,
                161,
                161,
                161,
                161,
                161,
                163,
                165,
                166,
                166,
                168,
                168,
                169,
                169,
                169,
                169,
                170,
                170,
                170,
                170,
                171,
                171,
                171,
                171,
                172,
                172,
                172,
                172,
                173,
                173,
                173,
                173,
                174,
                174,
                175,
                175,
                174,
                173,
                172,
                164,
                163,
                162,
                161,
                161,
                160,
                160,
                160,
                159,
                158,
                157,
                157,
                156,
                156,
                156,
                154,
                155,
                155,
                155,
                155,
                156,
                156,
                157,
                157,
                157,
                157,
                158,
                158,
                158,
                158,
                159,
                159,
                160,
                160,
                160,
                160,
                161,
                166,
                166,
                168,
                168,
                169,
                169,
                170,
                171,
                171,
                172,
                172,
                172,
                172,
                173,
                173,
                174,
                174,
                174,
                174,
                175,
                175,
                175,
                175,
                176,
                176,
                178,
                178,
                180,
                180,
                181,
                181,
                182,
                182,
                182,
                182,
                183,
                183,
                183,
                185,
                185,
                183,
                183,
                183,
                182,
                182,
                182,
                182,
                182,
                182,
                181,
                181,
                181,
                181,
                180,
                180,
                180,
                180,
                179,
                179,
                179,
                179,
                178,
                178,
                177,
                177,
                177,
                177,
                176,
                176,
                175,
                175,
                175,
                175,
                173,
                173,
                171,
                171,
                170,
                170,
                169,
                169,
                168,
                167,
                166,
                166,
                165,
                165,
                164,
                164,
                162,
                162,
                161,
                161,
                160,
                160,
                160,
                159,
                159,
                158,
                158,
                157,
                156,
                155,
                155,
                154,
                153,
                153,
                153,
                153,
                153,
                153,
                154,
                154,
                154,
                154,
                155,
                156,
                157,
                156,
                156,
                155,
                155,
                154,
                153,
                153,
                152,
                152,
                152,
                150,
                150,
            ],
            "y": [
                83,
                85,
                88,
                90,
                90,
                90,
                91,
                92,
                92,
                93,
                93,
                95,
                96,
                100,
                101,
                105,
                105,
                106,
                109,
                109,
                111,
                112,
                112,
                113,
                113,
                114,
                114,
                115,
                115,
                117,
                117,
                119,
                119,
                120,
                120,
                122,
                122,
                124,
                124,
                126,
                126,
                127,
                127,
                128,
                128,
                129,
                129,
                130,
                130,
                131,
                131,
                132,
                133,
                133,
                134,
                134,
                135,
                135,
                135,
                136,
                137,
                137,
                138,
                139,
                140,
                140,
                141,
                141,
                142,
                142,
                143,
                143,
                144,
                144,
                145,
                145,
                146,
                146,
                147,
                148,
                149,
                149,
                151,
                151,
                152,
                152,
                152,
                152,
                153,
                153,
                153,
                153,
                154,
                154,
                154,
                154,
                156,
                156,
                160,
                160,
                159,
                159,
                158,
                158,
                157,
                157,
                157,
                157,
                157,
                157,
                156,
                156,
                156,
                156,
                155,
                155,
                154,
                154,
                154,
                154,
                153,
                153,
                152,
                152,
                152,
                150,
                150,
                149,
                148,
                148,
                147,
                147,
                146,
                146,
                146,
                145,
                144,
                143,
                143,
                142,
                142,
                143,
                143,
                144,
                147,
                148,
                149,
                149,
                151,
                151,
                156,
                156,
                158,
                159,
                159,
                160,
                160,
                163,
                163,
                164,
                166,
                166,
                172,
                173,
                176,
                176,
                180,
                181,
                183,
                184,
                186,
                186,
                188,
                188,
                190,
                190,
                193,
                193,
                195,
                195,
                197,
                197,
                198,
                198,
                199,
                199,
                200,
                201,
                201,
                202,
                203,
                203,
                205,
                205,
                207,
                208,
                212,
                212,
                215,
                215,
                217,
                217,
                219,
                220,
                221,
                222,
                222,
                223,
                223,
                224,
                225,
                226,
                226,
                226,
                226,
                227,
                227,
                228,
                229,
                229,
                232,
                232,
                233,
                234,
                234,
                236,
                236,
                236,
                236,
                237,
                237,
                237,
                237,
                238,
                238,
                237,
                237,
                236,
                235,
                235,
                234,
                234,
                233,
                232,
                232,
                228,
                228,
                227,
                226,
                225,
                225,
                225,
                224,
                223,
                223,
                222,
                221,
                221,
                220,
                219,
                218,
                217,
                217,
                217,
                216,
                216,
                215,
                215,
                214,
                213,
                213,
                211,
                211,
                210,
                209,
                208,
                207,
                206,
                206,
                203,
                203,
                202,
                201,
                201,
                201,
                201,
                203,
                204,
                205,
                207,
                207,
                209,
                209,
                214,
                214,
                215,
                216,
                216,
                216,
                216,
                213,
                213,
                212,
                212,
                211,
                210,
                210,
                210,
                209,
                208,
                208,
                207,
                206,
                205,
                204,
                204,
                203,
                203,
                202,
                202,
                202,
                201,
                200,
                199,
                199,
                196,
                195,
                195,
                195,
                195,
                194,
                194,
                193,
                193,
                193,
                193,
                192,
                192,
                191,
                191,
                190,
                190,
                190,
                189,
                188,
                187,
                187,
                186,
                186,
                185,
                184,
                184,
                183,
                183,
                182,
                181,
                180,
                180,
                178,
                178,
                177,
                176,
                175,
                175,
                174,
                172,
                172,
                172,
                172,
                172,
                172,
                171,
                171,
                170,
                170,
                170,
                169,
                169,
                168,
                167,
                164,
                164,
                163,
                163,
                162,
                162,
                161,
                161,
                160,
                159,
                158,
                157,
                157,
                154,
                154,
                153,
                152,
                152,
                151,
                150,
                149,
                149,
                148,
                148,
                148,
                142,
                142,
                142,
                140,
                140,
                136,
                135,
                133,
                133,
                132,
                132,
                131,
                131,
                130,
                129,
                129,
                128,
                128,
                127,
                127,
                127,
                126,
                126,
                125,
                124,
                124,
                123,
                123,
                122,
                122,
                120,
                120,
                120,
                119,
                118,
                117,
                115,
                115,
                114,
                113,
                112,
                112,
                110,
                110,
                109,
                109,
                109,
                109,
                108,
                108,
                107,
                107,
                106,
                106,
                106,
                106,
                105,
                105,
                105,
                105,
                105,
                105,
                103,
                103,
                103,
                103,
                102,
                101,
                101,
                97,
                97,
                96,
                95,
                94,
                93,
                92,
                91,
                91,
                90,
                90,
                89,
                88,
                88,
                88,
                88,
                87,
                87,
                87,
                87,
                85,
                83,
            ],
        },
    ]


def test_yolo_pose_from_results_empty(running_img):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={},
        obb=torch.empty((0, 7)),
    )

    model = YoloPose.from_result(result)
    assert model.model_dump() == {
        "cls": -1,
        "name": "",
        "confidence": 0,
        "box": {
            "coords": [],
            "title": "",
        },
        "pose": {
            "x": [],
            "y": [],
            "visible": [],
        },
    }

    model = YoloPoses.from_results([result])
    assert model.model_dump() == {
        "cls": [],
        "name": [],
        "confidence": [],
        "box": [],
        "pose": [],
    }


def test_yolo_pose_from_results_empty_poses(running_img):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={0: "person"},
        boxes=torch.tensor([[102.0, 84.0, 183.0, 238.0, 0.9078, 0.0]]),
    )

    model = YoloPose.from_result(result)
    assert model.model_dump() == {
        "cls": 0,
        "name": "person",
        "confidence": 0.9078,
        "box": {
            "coords": [102, 84, 183, 238],
            "title": "person",
        },
        "pose": {
            "x": [],
            "y": [],
            "visible": [],
        },
    }

    model = YoloPoses.from_results([result])
    assert model.model_dump() == {
        "cls": [0],
        "name": ["person"],
        "confidence": [0.9078],
        "box": [{"coords": [102, 84, 183, 238], "title": "person"}],
        "pose": [],
    }


def test_yolo_pose_from_results(running_img, running_img_masks):
    result = Results(
        orig_img=running_img,
        path="running.jpeg",
        names={0: "person"},
        boxes=torch.tensor([[102.0, 84.0, 183.0, 238.0, 0.9078, 0.0]]),
        keypoints=torch.tensor(
            [
                [
                    [138.4593, 103.4513, 0.991],
                    [141.6461, 100.4834, 0.9863],
                    [136.6068, 100.2756, 0.8826],
                    [150.1754, 99.0366, 0.9691],
                    [0.0, 0.0, 0.2733],
                    [163.061, 114.3064, 0.9993],
                    [134.8462, 113.5232, 0.9977],
                    [179.6755, 136.7466, 0.9974],
                    [127.26, 130.8018, 0.9806],
                    [172.8256, 158.7771, 0.9962],
                    [118.0251, 146.8072, 0.9786],
                    [154.7233, 160.504, 0.9999],
                    [139.7511, 158.7988, 0.9997],
                    [143.9287, 193.5594, 0.9995],
                    [145.329, 190.7517, 0.9991],
                    [137.9563, 223.4, 0.9941],
                    [156.7822, 218.5813, 0.9921],
                ]
            ],
        ),
    )

    model = YoloPose.from_result(result)
    model_dict = model.model_dump()
    assert set(model_dict.keys()) == {
        "cls",
        "name",
        "confidence",
        "box",
        "pose",
    }
    assert model_dict["cls"] == 0
    assert model_dict["name"] == "person"
    assert_almost_equal(model_dict["confidence"], 0.9078, decimal=3)
    assert model_dict["box"] == {
        "coords": [102, 84, 183, 238],
        "title": "person",
    }
    assert set(model_dict["pose"]) == {"x", "y", "visible"}
    assert model_dict["pose"]["x"] == [
        138,
        142,
        137,
        150,
        0,
        163,
        135,
        180,
        127,
        173,
        118,
        155,
        140,
        144,
        145,
        138,
        157,
    ]
    assert model_dict["pose"]["y"] == [
        103,
        100,
        100,
        99,
        0,
        114,
        114,
        137,
        131,
        159,
        147,
        161,
        159,
        194,
        191,
        223,
        219,
    ]
    assert_array_almost_equal(
        model_dict["pose"]["visible"],
        [
            0.991,
            0.9862,
            0.8826,
            0.9691,
            0.2733,
            0.9993,
            0.9977,
            0.9974,
            0.9806,
            0.9962,
            0.9786,
            0.9999,
            0.9997,
            0.9995,
            0.9991,
            0.9941,
            0.9921,
        ],
        decimal=3,
    )

    model = YoloPoses.from_results([result])
    model_dict = model.model_dump()
    assert set(model_dict.keys()) == {
        "cls",
        "name",
        "confidence",
        "box",
        "pose",
    }
    assert model_dict["cls"] == [0]
    assert model_dict["name"] == ["person"]
    assert_array_almost_equal(model_dict["confidence"], [0.9078], decimal=3)
    assert model_dict["box"] == [{"coords": [102, 84, 183, 238], "title": "person"}]
    assert len(model_dict["pose"]) == 1
    assert set(model_dict["pose"][0]) == {"x", "y", "visible"}
    assert model_dict["pose"][0]["x"] == [
        138,
        142,
        137,
        150,
        0,
        163,
        135,
        180,
        127,
        173,
        118,
        155,
        140,
        144,
        145,
        138,
        157,
    ]
    assert model_dict["pose"][0]["y"] == [
        103,
        100,
        100,
        99,
        0,
        114,
        114,
        137,
        131,
        159,
        147,
        161,
        159,
        194,
        191,
        223,
        219,
    ]
    assert_array_almost_equal(
        model_dict["pose"][0]["visible"],
        [
            0.991,
            0.9862,
            0.8826,
            0.9691,
            0.2733,
            0.9993,
            0.9977,
            0.9974,
            0.9806,
            0.9962,
            0.9786,
            0.9999,
            0.9997,
            0.9995,
            0.9991,
            0.9941,
            0.9921,
        ],
        decimal=3,
    )
